import warnings
warnings.filterwarnings('ignore')
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
import tensorflow as tf
tf.compat.v1.logging.set_verbosity(40)

from tensorflow.keras import utils

train_Datadir = '../dl2_DataSets/cifar2/train/'
test_Datadir = '../dl2_DataSets/cifar2/test/'

import cv2
import numpy as np

def readData(path):
    x_images = []
    y_labels = []
    for i, j in enumerate(os.listdir(path)):
        sub_path = os.path.join(path, j)
        for image_name in os.listdir(sub_path):
            image_path = os.path.join(sub_path, image_name)
            image = cv2.imread(image_path)/255
            x_images.append(image)
            y_labels.append(i)
    return np.array(x_images), np.array(y_labels)

x_train, y_train = readData(train_Datadir)
x_test, y_test = readData(test_Datadir)

y_train = utils.to_categorical(y_train, 2)
y_test = utils.to_categorical(y_test, 2)

# data pipeline
db_train = tf.data.Dataset.from_tensor_slices((x_train, y_train)).shuffle(9000).batch(batch_size=128)
db_test = tf.data.Dataset.from_tensor_slices((x_test, y_test)).shuffle(1000).batch(batch_size=128)


class ConvCell(tf.keras.models.Model):
    def __init__(self, ch, kernize=3, stride=1):
        super(ConvCell, self).__init__()
        self.model = tf.keras.Sequential([
            tf.keras.layers.Conv2D(ch, kernize, strides=stride, padding='same'),
            tf.keras.layers.BatchNormalization(),
            tf.keras.layers.Activation('relu')
        ])
    def call(self, input):
        x = self.model(input, training=False)
        return x

class InceptionBlock(tf.keras.models.Model):
    def __init__(self, ch, stride=1):
        super(InceptionBlock, self).__init__()
        self.ch = ch
        self.stride = stride
        self.conv1 = ConvCell(ch, kernize=1, stride=stride)
        self.conv2_1 = ConvCell(ch, kernize=1, stride=stride)
        self.conv2_2 = ConvCell(ch, kernize=3, stride=1)
        self.conv3_1 = ConvCell(ch, kernize=1, stride=stride)
        self.conv3_2 = ConvCell(ch, kernize=5, stride=1)
        self.conv4_1 = tf.keras.layers.MaxPool2D(pool_size=3, strides=1, padding='same')
        self.conv4_2 = ConvCell(ch, kernize=1, stride=stride)

    def call(self,x):
        x1 = self.conv1(x)
        x2_1 = self.conv2_1(x)
        x2_2 = self.conv2_2(x2_1)
        x3_1 = self.conv3_1(x)
        x3_2 = self.conv3_2(x3_1)
        x4_1 = self.conv4_1(x)
        x4_2 = self.conv4_2(x4_1)
        x = tf.concat([x1, x2_2, x3_2, x4_2], axis=3)
        return x

class InceptionNet10(tf.keras.models.Model):
    def __init__(self, num_blocks, nb_class, init_ch=16):
        super(InceptionNet10, self).__init__()
        self.in_ch = init_ch
        self.out_ch = init_ch
        self.num_blocks = num_blocks

        self.conv1 = ConvCell(init_ch)
        self.blocks = tf.keras.Sequential()
        for blocks_id in range(num_blocks):
            for layers_id in range(2):
                if layers_id == 0:
                    block = InceptionBlock(self.out_ch, stride=2)
                else:
                    block = InceptionBlock(self.out_ch, stride=1)
                self.blocks.add(block)
            self.out_ch *= 2

        self.avgpool = tf.keras.layers.GlobalAveragePooling2D()
        self.fc = tf.keras.layers.Dense(nb_class, activation='softmax')

    def call(self, x):
        x = self.conv1(x)
        x = self.blocks(x)
        x = self.avgpool(x)
        x = self.fc(x)
        return x

if __name__ == '__main__':

    model = InceptionNet10(num_blocks=2, nb_class=2)

    model.compile(loss=tf.keras.losses.CategoricalCrossentropy(),
                  optimizer=tf.keras.optimizers.Adam(),
                  metrics=['accuracy'])

    history = model.fit(db_train, validation_data=db_test, epochs=6)

    acc = history.history['accuracy']
    val_acc = history.history['val_accuracy']

    loss = history.history['loss']
    val_loss = history.history['val_loss']

    import matplotlib.pyplot as plt

    plt.plot(acc, label='acc')
    plt.plot(val_acc, label='val_acc')
    plt.legend()
    plt.show()

    plt.plot(loss, label='loss')
    plt.plot(val_loss, label='val_loss')
    plt.legend()
    plt.show()
